import requests
import json
import time
import os
from tencentcloud.common import credential
from tencentcloud.hunyuan.v20230901 import hunyuan_client, models as hunyuan_models

class AIModelConnector:
    def __init__(self, config):
        self.config = config
        self.model_cache = {}
        
    def analyze_with_hunyuan(self, data, prompt_template):
        """使用腾讯混元大模型进行分析"""
        # 配置腾讯云API凭证
        cred = credential.Credential(
            self.config['tencent_secret_id'],
            self.config['tencent_secret_key']
        )
        
        # 创建客户端对象
        client = hunyuan_client.HunyuanClient(cred, "ap-guangzhou")
        
        # 组装请求数据
        formatted_prompt = prompt_template.format(stock_data=json.dumps(data, ensure_ascii=False))
        
        # 创建请求对象
        req = hunyuan_models.ChatCompletionsRequest()
        req.Messages = [{"role": "user", "content": formatted_prompt}]
        req.Model = "hunyuan-standard"
        
        # 调用API
        resp = client.ChatCompletions(req)
        
        return json.loads(resp.Choices[0].Message.Content)
    
    def analyze_with_deepseek(self, data, prompt_template):
        """使用DeepSeek模型进行分析"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.config['deepseek_api_key']}"
        }
        
        formatted_prompt = prompt_template.format(stock_data=json.dumps(data, ensure_ascii=False))
        
        payload = {
            "model": "deepseek-chat",
            "messages": [{"role": "user", "content": formatted_prompt}],
            "temperature": 0.1,
            "max_tokens": 2000
        }
        
        response = requests.post(
            "https://api.deepseek.com/v1/chat/completions",
            headers=headers,
            data=json.dumps(payload)
        )
        
        return json.loads(response.json()['choices'][0]['message']['content'])

    def get_analysis_result(self, stock_code, stock_data, news_data, analysis_type):
        """根据分析类型选择合适的模型进行分析"""
        if analysis_type == "summary":
            # 使用腾讯混元进行总结分析
            prompt = self.config['prompt_templates']['summary']
            return self.analyze_with_hunyuan({
                'stock_code': stock_code,
                'stock_data': stock_data,
                'news_data': news_data
            }, prompt)
        
        elif analysis_type == "technical":
            # 使用DeepSeek进行技术面分析
            prompt = self.config['prompt_templates']['technical']
            return self.analyze_with_deepseek({
                'stock_code': stock_code,
                'stock_data': stock_data
            }, prompt)
        
        elif analysis_type == "fundamental":
            # 使用DeepSeek进行基本面分析
            prompt = self.config['prompt_templates']['fundamental']
            return self.analyze_with_deepseek({
                'stock_code': stock_code,
                'stock_data': stock_data,
                'news_data': news_data
            }, prompt)

    def test_connection(self):
        """测试AI模型连接是否正常"""
        try:
            # 简单测试
            test_prompt = "这是一个测试"
            if 'tencent_secret_id' in self.config and 'tencent_secret_key' in self.config:
                # 测试腾讯云连接
                return True
            elif 'deepseek_api_key' in self.config:
                # 测试DeepSeek连接
                return True
            return False
        except Exception as e:
            logger.error(f"AI模型连接测试失败: {str(e)}")
            return False 